# encoding: utf-8
"""
@author:  sherlock
@contact: sherlockliao01@gmail.com
"""

import argparse
import os
import sys
from os import mkdir

from PIL import Image
import torch
from torch.backends import cudnn

sys.path.append('.')
from configa import cfg
from data import make_data_loader
from modeling import build_model
from utils.logger import setup_logger
import logging




def inference(model,data):
    model.eval()
    with torch.no_grad():
        # add by fby
        data = data.to("cuda")
        feat = model(data)

        # add by fby
        return feat

def get_feats(model ,dataset, cfg,dim):
    f = torch.randn(len(dataset),dim)
    for i in range(len(dataset)):
        data, _ = dataset[i]
        feat = inference(model, data)
        if cfg.TEST.FEAT_NORM == 'yes':
            feat = torch.nn.functional.normalize(feat, dim=1, p=2)
        f[i] = feat
    return f

def main():
    parser = argparse.ArgumentParser(description="ReID Baseline Inference")
    parser.add_argument(
        "--config_file", default="/home/fby/reid-strong-baseline/app/fby.yml", help="path to config file", type=str
    )
    parser.add_argument("opts", help="Modify config options using the command-line", default=None,
                        nargs=argparse.REMAINDER)

    args = parser.parse_args()

    num_gpus = int(os.environ["WORLD_SIZE"]) if "WORLD_SIZE" in os.environ else 1

    if args.config_file != "":
        cfg.merge_from_file(args.config_file)
    cfg.merge_from_list(args.opts)
    cfg.freeze()

    output_dir = cfg.OUTPUT_DIR
    if output_dir and not os.path.exists(output_dir):
        mkdir(output_dir)

    logger = setup_logger("reid_baseline", os.path.dirname(os.path.abspath(cfg["TEST"]["WEIGHT"])), 0, fn="test.txt")
    logger.info("Using {} GPUS".format(num_gpus))
    logger.info(args)

    if args.config_file != "":
        logger.info("Loaded configuration file {}".format(args.config_file))
        with open(args.config_file, 'r') as cf:
            config_str = "\n" + cf.read()
            logger.info(config_str)
    logger.info("Running with config:\n{}".format(cfg))

    if cfg.MODEL.DEVICE == "cuda":
        os.environ['CUDA_VISIBLE_DEVICES'] = cfg.MODEL.DEVICE_ID
    cudnn.benchmark = True

    num_classes = 483
    model = build_model(cfg, num_classes)
    dim = model.in_planes
    model.load_param(cfg.TEST.WEIGHT)
    model.cuda()

    # inference(cfg, model, num_query)

    from app_dataset_loader import ImageDataset

    g_dataset = ImageDataset(cfg.TEST.DATA_DIR_g,cfg)

    gf = get_feats(model,g_dataset,cfg,dim)

    q_dataset = ImageDataset(cfg.TEST.DATA_DIR_q,cfg)
    qf = get_feats(model, q_dataset, cfg,dim)

    m, n = qf.shape[0], gf.shape[0]
    distmat = torch.pow(qf, 2).sum(dim=1, keepdim=True).expand(m, n) + \
              torch.pow(gf, 2).sum(dim=1, keepdim=True).expand(n, m).t()
    distmat.addmm_(1, -2, qf, gf.t())
    distmat = distmat.cpu().numpy()

    import numpy as np
    order_index = np.argsort(distmat, axis=1)
    save_dir = '/home/fby/MobileNet-YOLOv5/runs/detect/exp4/crops/2'
    for i in range(1):
        # 找j 个最像的
        for j in range(0, 30):
            index = order_index[i][j]
            _, gallery_name = g_dataset[index]
            gallery_fn = os.path.join(save_dir, str(j)+"__"+ gallery_name.split('/')[-1])
            gallery_im = Image.open(gallery_name)
            gallery_im.save(gallery_fn)

if __name__ == '__main__':
    main()
